import os
import csv
import random
from glob import glob

def generate_csv_from_folder(root_folder, output_csv):
    # 定义疾病分类的文件夹和对应的列名
    disease_categories = {
        'normal': ('N', 'normal fundus'),
        'diabetes': ('D', 'moderate non proliferative retinopathy'),
        'glaucoma': ('G', 'glaucoma'),
        'cataract': ('C', 'cataract'),
        'amd': ('A', 'amd'),
        'hypertension': ('H', 'hypertensive retinopathy'),
        'myopia': ('M', 'macular epiretinal membrane'),
        'other': ('O', 'other')  # 包括drusen, laser spot等
    }
    
    # 准备写入CSV文件
    with open(output_csv, mode='w', newline='', encoding='utf-8') as csv_file:
        writer = csv.writer(csv_file)
        
        # 写入表头
        header = ['ID', 'Patient Age', 'Patient Sex', 'Fundus', 'Diagnostic Keywords', 
                 'N', 'D', 'G', 'C', 'A', 'H', 'M', 'O']
        writer.writerow(header)
        
        id_counter = 0
        
        # 遍历每个疾病分类文件夹
        for folder_name, (category, keyword) in disease_categories.items():
            category_folder = os.path.join(root_folder, folder_name)
            
            if not os.path.exists(category_folder):
                print(f"Warning: Folder {category_folder} does not exist. Skipping...")
                continue
                
            # 获取该文件夹下所有图像文件
            image_files = glob(os.path.join(category_folder, '*.jpg')) + \
                        glob(os.path.join(category_folder, '*.jpeg')) + \
                        glob(os.path.join(category_folder, '*.png'))
            
            # 为每个图像文件生成一行数据
            for image_path in image_files:
                # 获取图像文件名
                image_name = os.path.basename(image_path)
                fundus_path = f"{folder_name}/{image_name}"  # 子文件夹名/图像名
                
                # 随机生成年龄和性别
                age = random.randint(20, 80)
                sex = random.choice(['Male', 'Female'])
                
                # 初始化所有疾病分类为0
                disease_values = {c: 0 for c in ['N', 'D', 'G', 'C', 'A', 'H', 'M', 'O']}
                # 设置当前分类为1
                disease_values[category] = 1
                
                # 准备行数据
                row = [
                    id_counter,
                    age,
                    sex,
                    fundus_path,
                    keyword,
                    disease_values['N'],
                    disease_values['D'],
                    disease_values['G'],
                    disease_values['C'],
                    disease_values['A'],
                    disease_values['H'],
                    disease_values['M'],
                    disease_values['O']
                ]
                
                writer.writerow(row)
                id_counter += 1

if __name__ == "__main__":
    # 设置你的文件夹路径和输出CSV文件名
    input_folder = "/data/zhangyichi/collected_dataset/"  # 替换为你的文件夹路径
    output_csv = "/home/zhangyichi/dataset/OIA-ODIR/Training Set/AOD.csv"  # 输出CSV文件名
    
    # 确保文件夹结构正确
    # 你的文件夹应该包含子文件夹N, D, G, C, A, H, M, O
    # 每个子文件夹中包含对应分类的图像
    
    generate_csv_from_folder(input_folder, output_csv)
    print(f"CSV文件已生成: {output_csv}")